"""
optimizer_config.py
-------------------
定义优化器与学习率调度器（Scheduler）相关的配置。
这些参数用于训练器中初始化优化器和调度策略。

支持的优化器包括：
- SGD, Adam, AdamW, RMSprop, Adagrad, LAMB

支持的调度器包括：
- StepLR, MultiStepLR, CosineAnnealingLR, OneCycleLR, ReduceLROnPlateau
"""

from dataclasses import dataclass, field
from typing import Optional, List, Dict, Union, Tuple


# ======================
# 优化器配置
# ======================
@dataclass
class OptimizerConfig:
    optimizer_type: str = "adam"                    # 优化器类型
    learning_rate: float = 1e-3                     # 学习率
    weight_decay: float = 1e-4                      # 权重衰减
    momentum: float = 0.9                           # SGD 动量
    betas: tuple[Union[float], Union[float]] = (0.9, 0.999)     # Adam / AdamW 参数
    eps: float = 1e-8                               # 数值稳定参数
    amsgrad: bool = False                           # Adam 的 AMSGrad 选项
    rho: float = 0.9                                # RMSprop 衰减率
    alpha: float = 0.99                             # RMSprop 平滑参数
    centered: bool = False                          # RMSprop 中心化
    clip_grad_norm: Optional[float] = None          # 梯度裁剪阈值

    def summary(self):
        """打印优化器配置"""
        print("⚙️ 优化器配置")
        print(f"  类型: {self.optimizer_type}")
        print(f"  学习率: {self.learning_rate}, 权重衰减: {self.weight_decay}")
        print(f"  动量: {self.momentum}, betas: {self.betas}, amsgrad: {self.amsgrad}")
        if self.clip_grad_norm:
            print(f"  梯度裁剪阈值: {self.clip_grad_norm}")


# ======================
# 调度器配置
# ======================
@dataclass
class SchedulerConfig:
    scheduler_type: Optional[str] = None
    """
    调度器类型，可选：
    - "steplr"
    - "multisteplr"
    - "cosine"
    - "onecycle"
    - "plateau"
    """

    step_size: int = 30                # StepLR 步长
    gamma: float = 0.1                 # 学习率衰减系数
    milestones: Optional[List[int]] = None  # MultiStepLR 的关键点
    min_lr: float = 1e-6               # 最小学习率（CosineAnnealing）
    max_lr: Optional[float] = None     # 最大学习率（OneCycle）
    warmup_epochs: int = 0             # 预热期 epoch 数
    patience: int = 10                 # ReduceLROnPlateau 容忍次数
    mode: str = "min"                  # ReduceLROnPlateau 模式（min/max）
    factor: float = 0.1                # ReduceLROnPlateau 衰减因子
    last_epoch: int = -1               # 上一 epoch 索引

    def summary(self):
        """打印调度器配置"""
        print("📉 学习率调度器配置")
        if not self.scheduler_type:
            print("  未启用调度器")
            return
        print(f"  类型: {self.scheduler_type}")
        print(f"  参数: step_size={self.step_size}, gamma={self.gamma}, min_lr={self.min_lr}")


# ======================
# 总配置封装
# ======================
@dataclass
class OptimizerManagerConfig:
    optimizer: OptimizerConfig = None
    scheduler: SchedulerConfig = None

    def summary(self):
        print("============== 优化器与调度器配置 ==============")
        self.optimizer.summary()
        self.scheduler.summary()
        print("================================================")
